import datetime
import os
import shutil
from turtledemo.chaos import line
from typing import Tuple
import fitz  # PyMuPDF
import cv2
from paddleocr import PPStructure, save_structure_res
from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx


class PDFProcessor:
    def __init__(self, dpi: int = 150, lang: str = 'ch'):
        # 初始化dpi属性，默认为600
        self.dpi = dpi
        # 初始化lang属性，默认为'ch'
        self.lang = lang
        # 初始化table_engine属性，使用PPStructure类进行表格和OCR处理
        # 启用表格处理功能
        # 启用OCR功能
        # 使用传入的lang作为OCR和表格处理的语言
        # 不显示日志信息
        self.table_engine = PPStructure(table=True, ocr=True, lang=lang, show_log=False)


    def _create_dir(self, path: str) -> None:
        """创建目录并清空已有内容"""
        # 检查指定路径是否存在
        if os.path.exists(path):
            # 如果路径存在，则删除该路径下的所有文件和子目录
            shutil.rmtree(path)
        # 创建指定路径的目录，如果目录已存在则不抛出异常
        os.makedirs(path, exist_ok=True)


    def pdf_to_images(self, pdf_path: str, output_dir: str) -> Tuple[str, int]:
        """将PDF转换为高质量PNG图像"""
        # 记录开始时间
        start_time = datetime.datetime.now()

        # 检查PDF文件是否存在
        if not os.path.isfile(pdf_path):
            raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")

        # 创建干净输出目录
        # 获取PDF文件名（不含扩展名）
        base_name = os.path.splitext(os.path.basename(pdf_path))[0]
        # 构建输出目录路径
        image_dir = os.path.join(output_dir, base_name)
        # 调用内部方法创建或清空目录
        self._create_dir(image_dir)

        # 使用fitz库打开PDF文件
        # noinspection PyUnresolvedReferences
        with fitz.open(pdf_path) as pdf_doc:
            # 获取PDF总页数
            total_pages = len(pdf_doc)
            # 根据DPI计算缩放比例
            zoom = self.dpi / 72  # 根据DPI计算缩放比例
            # 创建缩放矩阵
            matrix = fitz.Matrix(zoom, zoom)

            # 遍历PDF每一页
            for page_num in range(total_pages):
                # 加载当前页
                page = pdf_doc.load_page(page_num)
                # 使用缩放矩阵获取当前页的像素图
                pix = page.get_pixmap(matrix=matrix, alpha=False)
                # 构建输出图像路径
                img_path = os.path.join(image_dir, f"page_{page_num + 1:03d}.png")
                # 保存像素图为PNG图像
                pix.save(img_path)
                # 打印转换进度
                print(f"已转换 {os.path.basename(pdf_path)} 第 {page_num + 1}/{total_pages} 页")

        # 计算并打印转换耗时
        time_cost = (datetime.datetime.now() - start_time).seconds
        print(f"PDF转换完成，共 {total_pages} 页，耗时 {time_cost} 秒")

        # 返回输出目录和总页数
        return image_dir, total_pages


    def process_images(self, image_dir: str, output_dir: str) -> None:
        """处理图像并生成Word文档"""
        # 记录开始时间
        start_time = datetime.datetime.now()
        # 创建输出目录
        self._create_dir(output_dir)

        # 获取排序后的图像文件列表
        # 按文件名（假设格式为'page_XX.png'）排序
        image_files = sorted(
            # 过滤出以'.png'结尾的文件
            [f for f in os.listdir(image_dir) if f.lower().endswith('.png')],
            # 使用文件名中的数字进行排序
            key=lambda x: int(x.split('_')[1].split('.')[0])
        )

        # 存储完整文本内容的列表
        full_text = []
        # 遍历排序后的图像文件列表
        for idx, img_file in enumerate(image_files, 1):
            img_path = os.path.join(image_dir, img_file)
            print(f"\n正在处理第 {idx} 页: {img_file}")

            try:
                # 读取图像文件
                img = cv2.imread(img_path)
                # 检查图像是否成功读取
                if img is None:
                    raise IOError(f"无法读取图像文件: {img_path}")

                # 执行版面分析
                result = self.table_engine(img)
                # 获取图像的高度和宽度
                h, w = img.shape[:2]

                # 保存分析结果
                save_structure_res(result, output_dir, f"page_{idx:03d}")

                # 排序并转换文档
                sorted_result = sorted_layout_boxes(result, w)
                convert_info_docx(img, sorted_result, output_dir, f"page_{idx:03d}")

                # 收集文本内容
                for item in sorted_result:
                    for para in item['res']:
                        full_text.append(para['text'])
                    full_text.append('\n')

            except Exception as e:
                # 打印错误信息并继续处理下一张图像
                print(f"处理第 {idx} 页时发生错误: {str(e)}")
                continue

        # 保存完整文本到文件
        txt_path = os.path.join(output_dir, 'full_text.txt')
        with open(txt_path, 'w', encoding='utf-8') as f:
            f.write(''.join(map(lambda x:x[-1],line)))


        # 计算并打印处理耗时
        time_cost = (datetime.datetime.now() - start_time).seconds
        print(f"\n文档处理完成，耗时 {time_cost} 秒")
        print(f"结果已保存至: {output_dir}")


if __name__ == "__main__":
    PDF_PATH= r'./data/2.pdf'
    BASE_DIR = './output'
    IMAGE_DIR = './imgs/2'
    IMAGE_DIR = os.path.join(BASE_DIR, "images")
    RESULT_DIR = os.path.join(BASE_DIR, "results")

    # 初始化处理器
    processor = PDFProcessor(dpi=150, lang='ch')

    try:
        # 第一步：PDF转图片
        image_folder, total_pages = processor.pdf_to_images(PDF_PATH, IMAGE_DIR)

        # 第二步：处理图片生成文档
        processor.process_images(image_folder, RESULT_DIR)

    except Exception as e:
        print(f"\n处理过程中发生严重错误: {str(e)}")
    finally:
        # 可选：处理完成后删除临时图片
        # shutil.rmtree(IMAGE_DIR)
        pass


